import os
import sys
import logging
import time
from pyngrok import ngrok, conf

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('ngrok_connector')

def start_ngrok(port=5000, max_retries=3):
    """启动ngrok隧道并打印公网URL"""
    # 获取环境变量中的NGROK_AUTH_TOKEN，如果有的话
    auth_token = os.environ.get("NGROK_AUTH_TOKEN")
    if auth_token:
        logger.info("使用环境变量中的NGROK_AUTH_TOKEN")
        conf.get_default().auth_token = auth_token
    
    retry_count = 0
    while retry_count < max_retries:
        try:
            # 先尝试关闭任何已存在的ngrok连接
            tunnels = ngrok.get_tunnels()
            for tunnel in tunnels:
                logger.info(f"关闭现有隧道: {tunnel.public_url}")
                ngrok.disconnect(tunnel.public_url)
            
            # 等待一会儿让服务器有时间处理关闭请求
            if tunnels:
                logger.info("等待5秒让服务器处理关闭请求...")
                time.sleep(5)
            
            # 开启新的http隧道
            logger.info(f"正在为端口 {port} 创建ngrok隧道...")
            public_url = ngrok.connect(port).public_url
            
            logger.info(f"ngrok隧道创建成功！")
            logger.info(f"公网访问地址: {public_url}")
            logger.info(f"Web界面: {public_url}")
            logger.info(f"API提取对话示例: {public_url}/extract?url=YOUR_IMAGE_URL")
            logger.info(f"API状态检查: {public_url}/status")
            logger.info("按Ctrl+C停止服务")
            
            # 持续监控隧道状态
            while True:
                time.sleep(10)  # 每10秒检查一次
                try:
                    # 验证隧道是否仍然存在
                    current_tunnels = ngrok.get_tunnels()
                    if not current_tunnels:
                        logger.warning("隧道已关闭，正在尝试重新连接...")
                        public_url = ngrok.connect(port).public_url
                        logger.info(f"隧道重新连接成功: {public_url}")
                except Exception as e:
                    logger.error(f"检查隧道状态时出错: {e}")
                    # 尝试重新连接
                    try:
                        public_url = ngrok.connect(port).public_url
                        logger.info(f"隧道重新连接成功: {public_url}")
                    except:
                        logger.error("无法重新连接隧道")
            
        except Exception as e:
            retry_count += 1
            logger.error(f"启动ngrok隧道时出错 (尝试 {retry_count}/{max_retries}): {e}")
            if "account is limited to 1 simultaneous ngrok agent" in str(e):
                logger.error("您的ngrok账户限制只能运行一个隧道，尝试关闭所有隧道后重试")
                # 尝试关闭所有隧道
                try:
                    tunnels = ngrok.get_tunnels()
                    for tunnel in tunnels:
                        ngrok.disconnect(tunnel.public_url)
                    logger.info("已关闭所有现有隧道")
                except:
                    logger.error("无法关闭现有隧道")
                
                # 等待一会儿让服务器有时间处理关闭请求
                logger.info("等待10秒让服务器处理关闭请求...")
                time.sleep(10)
            else:
                # 其他错误，等待一段时间后重试
                wait_time = 5 * retry_count
                logger.info(f"等待 {wait_time} 秒后重试...")
                time.sleep(wait_time)
    
    # 如果达到最大重试次数仍然失败
    if retry_count >= max_retries:
        logger.error(f"达到最大重试次数 {max_retries}，无法启动ngrok隧道")
        logger.info("建议：")
        logger.info("1. 访问 https://dashboard.ngrok.com/agents 手动终止所有活跃会话")
        logger.info("2. 确认您的ngrok在本地没有其他实例在运行")
        logger.info("3. 检查您的网络连接")
        sys.exit(1)

if __name__ == "__main__":
    # 默认端口5000，可以通过命令行参数修改
    port = 5000
    if len(sys.argv) > 1:
        try:
            port = int(sys.argv[1])
        except ValueError:
            logger.error(f"无效的端口号: {sys.argv[1]}")
            sys.exit(1)
    
    start_ngrok(port)
